from django.db.models import Count

from api_app.data_model_manager.enums import DataModelEvaluations
from api_app.engines_manager.engines.evaluation import EvaluationEngineModule


class MalwareFamilyEngineModule(EvaluationEngineModule):
    def run(self):
        result = super().run()
        if result["evaluation"] in [
            DataModelEvaluations.TRUSTED.value,
        ]:
            first_family = None
        else:
            first_obj = (
                self.job.get_analyzers_data_models()
                .values("malware_family")
                .order_by("-malware_family")
                .annotate(count=Count("malware_family", distinct=True))
                .first()
            )
            if first_obj:
                first_family = first_obj["malware_family"]
            else:
                first_family = None

        return {
            **result,
            "malware_family": first_family,
        }
